import torch
from transformers import BertTokenizer

from emotionAnalysis.Net import Model

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(DEVICE)

values=["负向评价","正向评价"]
model=Model().to(DEVICE)

# 加载分词器
model_name = "./model/google-bert/bert-base-chinese/models--google-bert--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
token = BertTokenizer.from_pretrained(model_name)

# 自定义数据编码处理函数
def collate_fn(data):
    sente = [data]
    # 编码处理
    data = token.batch_encode_plus(
        batch_text_or_text_pairs=sente,
        truncation=True,
        padding='max_length',
        max_length=300,
        return_tensors='pt',
        return_length=True
    )
    input_ids = data['input_ids']
    attention_mask = data['attention_mask']
    token_type_ids = data['token_type_ids']

    return input_ids, attention_mask, token_type_ids

def test():
    model.load_state_dict(torch.load("./params/1bert.pt"))
    model.eval()
    while True:
        data=input("请输入测试数据（输入‘q’退出）:")
        if data=="q":
            print("测试结束")
            break
        input_ids, attention_mask, token_type_ids=collate_fn(data)
        input_ids, attention_mask, token_type_ids=input_ids.to(DEVICE), attention_mask.to(DEVICE), token_type_ids.to(DEVICE)

        with torch.no_grad():
            out=model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
            out=out.argmax(dim=1)
            print("模型判定：",values[out],"\n")

if __name__=="__main__":
    test()